
import cv2
import matplotlib.pyplot as plt
import numpy as np

def filp_augment(img,show=False):
    '''接受cv2图像，并进行水平翻转，返回翻转后图片(ndarrg)'''
    flipped_img = cv2.flip(img, 1)
    if show:
        plt.subplot(1, 2, 1)
        plt.imshow(img)
        plt.title('Original Image')
        plt.subplot(1, 2, 2)
        plt.imshow(flipped_img)
        plt.title('Flipped Image')
        plt.show()
    return flipped_img

def random_rotation(img,range:int = 30,show=False):
    '''接受cv2图像，并进行随机旋转，返回旋转后图片(ndarrg)'''
    angle = np.random.uniform(-range, range)

    height, width = img.shape[:2]
    center = (width // 2, height // 2)
    # 图像旋转
    rotation_matrix = cv2.getRotationMatrix2D(center, angle, 1.0)
    rotated_img = cv2.warpAffine(img, rotation_matrix, (width, height))
    # 显示图像
    if show:
        plt.subplot(1, 2, 1)
        plt.imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
        plt.title('Original Image')
        plt.subplot(1, 2, 2)
        plt.imshow(cv2.cvtColor(rotated_img, cv2.COLOR_BGR2RGB))
        plt.title('Rotated Image')
        plt.show()
    return rotated_img
def random_color(img, color_range = 30,bright_range = None,satura_range = None, show = False):
    '''接受cv2图像，并进行随机改变色调、亮度、饱和度，返回改变后图片(ndarrg)'''
    hsv_img = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
    if color_range:
        hue_offset = np.random.randint(-color_range, color_range)
        hsv_img[..., 0] = (hsv_img[..., 0] + hue_offset) % 180
    if bright_range:
        brightness_offset = np.random.randint(-bright_range, bright_range)
        hsv_img[..., 2] = np.clip(hsv_img[..., 2] + brightness_offset, 0, 255)
    if satura_range:
        saturation_offset = np.random.randint(-satura_range, satura_range)
        hsv_img[..., 1] = np.clip(hsv_img[..., 1] + saturation_offset, 0, 255)
    augmented_img = cv2.cvtColor(hsv_img, cv2.COLOR_HSV2BGR)

    if show:
        plt.subplot(1, 2, 1)
        plt.imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
        plt.title('Original Image')
        plt.subplot(1, 2, 2)
        plt.imshow(cv2.cvtColor(augmented_img, cv2.COLOR_BGR2RGB))
        plt.title('Augmented Image')
        plt.show()

    return augmented_img

def random_crop(img, show=False, crop_size=(224, 224)):
    '''接受cv2图像，并进行随机裁剪，返回裁剪后图片(ndarrg)'''
    height, width = img.shape[:2]

    crop_height, crop_width = crop_size
    max_x = width - crop_width
    max_y = height - crop_height

    x = np.random.randint(0, max_x + 1)
    y = np.random.randint(0, max_y + 1)

    cropped_img = img[y:y+crop_height, x:x+crop_width]

    if show:
        plt.subplot(1, 2, 1)
        plt.imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
        plt.title('Original Image')
        plt.subplot(1, 2, 2)
        plt.imshow(cv2.cvtColor(cropped_img, cv2.COLOR_BGR2RGB))
        plt.title('Cropped Image')
        plt.show()
    return cropped_img
def random_scale(img, show=False, scale_range=(0.5, 1.5)):
    '''接受cv2图像，并进行随机缩放，返回缩放后图片(ndarrg)'''
    scale_factor = np.random.uniform(scale_range[0], scale_range[1])

    scaled_img = cv2.resize(img, None, fx=scale_factor, fy=scale_factor, interpolation=cv2.INTER_LINEAR)

    if show:
        plt.subplot(1, 2, 1)
        plt.imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
        plt.title('Original Image')
        plt.subplot(1, 2, 2)
        plt.imshow(cv2.cvtColor(scaled_img, cv2.COLOR_BGR2RGB))
        plt.title('Scaled Image')
        plt.show()

    return scaled_img

if __name__ == '__main__':
    path = r'E:\A_project\Garbage_classify\datas\Garbage classification\plastic\plastic9.jpg'
    img = cv2.imread(path)
    img_ = random_crop(img,show=True)
    img_ = random_color(img,show=True)
    print(1)